from typing import List

import matplotlib.pyplot as plt
import numpy as np

from entity.task import Task
from utils.compute import calc_execute_time


def plot_gantt(tasks: List[Task], s_map: dict, p_map: dict, title: str, out_name: str):
    x = []
    y = []
    timeout_x = []
    timeout_y = []
    left = []
    timeout_left = []
    m_map = dict()
    for pid in p_map:
        m_map[pid] = 0
    for t in tasks:
        tid = t.id
        pid = s_map[tid]
        p = p_map[pid]
        exec_time = calc_execute_time(t, p)
        m = m_map[pid]
        m_map[pid] += exec_time

        if m_map[pid] > t.deadline:
            timeout_x.append(exec_time)
            timeout_y.append(pid)
            timeout_left.append(m)
        else:
            x.append(exec_time)
            y.append(pid)
            left.append(m)

    ticks = np.arange(1, 10)
    y_data = np.array(y)
    x_data = np.array(x)
    timeout_y_data = np.array(timeout_y)
    timeout_x_data = np.array(timeout_x)
    fig, ax = plt.subplots(1, 1)

    ax.barh(y_data, x_data, left=left, color='wheat',
            edgecolor='grey', label='scheduled')
    if len(timeout_y_data) > 0:
        ax.barh(timeout_y_data, timeout_x_data, left=timeout_left,
                color='salmon', edgecolor='grey', label='timeout')
    plt.legend(loc='lower right')
    plt.yticks(ticks)
    [ax.spines[i].set_visible(False) for i in ['top', 'right']]
    ax.set_title(title)
    ax.set_xlabel("length")
    ax.grid(alpha=0.5, axis='x')
    plt.savefig(f'{out_name}.png')
